import logging
import os
import re
import shutil
import socket
import sys
import tempfile
import threading
import time
from abc import ABC
from contextlib import contextmanager
from unittest import mock

import pytest
from pyspark.sql import SparkSession

import ray
import ray.util.spark.cluster_init
from ray._common.test_utils import wait_for_condition
from ray.util.spark import (
    MAX_NUM_WORKER_NODES,
    setup_global_ray_cluster,
    setup_ray_cluster,
    shutdown_ray_cluster,
)
from ray.util.spark.utils import (
    _calc_mem_per_ray_worker_node,
)

pytestmark = [
    pytest.mark.skipif(
        os.name != "posix",
        reason="Ray on spark only supports running on POSIX system.",
    ),
    pytest.mark.timeout(1500),
]


_RAY_ON_SPARK_WORKER_SHARED_MEMORY_BYTES = 2000000000
_RAY_ON_SPARK_WORKER_PHYSICAL_MEMORY_BYTES = 10000000000


def _setup_ray_on_spark_envs():
    os.environ["RAY_ON_SPARK_WORKER_SHARED_MEMORY_BYTES"] = str(
        _RAY_ON_SPARK_WORKER_SHARED_MEMORY_BYTES
    )
    os.environ["RAY_ON_SPARK_WORKER_PHYSICAL_MEMORY_BYTES"] = str(
        _RAY_ON_SPARK_WORKER_PHYSICAL_MEMORY_BYTES
    )


def setup_module():
    _setup_ray_on_spark_envs()


@contextmanager
def _setup_ray_cluster(*args, **kwds):
    setup_ray_cluster(*args, **kwds)
    try:
        yield ray.util.spark.cluster_init._active_ray_cluster
    finally:
        shutdown_ray_cluster()


_logger = logging.getLogger(__name__)


class RayOnSparkCPUClusterTestBase(ABC):
    spark = None
    num_total_cpus = None
    num_total_gpus = None
    num_cpus_per_spark_task = None
    num_gpus_per_spark_task = None
    max_spark_tasks = None

    @classmethod
    def teardown_class(cls):
        time.sleep(10)  # Wait all background spark job canceled.
        os.environ.pop("SPARK_WORKER_CORES", None)
        cls.spark.stop()

    @staticmethod
    def get_ray_worker_resources_list():
        wr_list = []
        for node in ray.nodes():
            # exclude dead node and head node (with 0 CPU resource)
            if node["Alive"] and node["Resources"].get("CPU", 0) > 0:
                wr_list.append(node["Resources"])
        return wr_list

    def test_cpu_allocation(self):
        for max_worker_nodes, num_cpus_worker_node, max_worker_nodes_arg in [
            (
                self.max_spark_tasks // 2,
                self.num_cpus_per_spark_task,
                self.max_spark_tasks // 2,
            ),
            (self.max_spark_tasks, self.num_cpus_per_spark_task, MAX_NUM_WORKER_NODES),
            (
                self.max_spark_tasks // 2,
                self.num_cpus_per_spark_task * 2,
                MAX_NUM_WORKER_NODES,
            ),
            (
                self.max_spark_tasks // 2,
                self.num_cpus_per_spark_task * 2,
                self.max_spark_tasks // 2 + 1,
            ),  # Test case: requesting resources exceeding all cluster resources
        ]:
            num_ray_task_slots = self.max_spark_tasks // (
                num_cpus_worker_node // self.num_cpus_per_spark_task
            )
            (
                mem_worker_node,
                object_store_mem_worker_node,
                _,
            ) = _calc_mem_per_ray_worker_node(
                num_task_slots=num_ray_task_slots,
                physical_mem_bytes=_RAY_ON_SPARK_WORKER_PHYSICAL_MEMORY_BYTES,
                shared_mem_bytes=_RAY_ON_SPARK_WORKER_SHARED_MEMORY_BYTES,
                configured_heap_memory_bytes=None,
                configured_object_store_bytes=None,
            )
            with _setup_ray_cluster(
                max_worker_nodes=max_worker_nodes_arg,
                num_cpus_worker_node=num_cpus_worker_node,
                num_gpus_worker_node=0,
                head_node_options={"include_dashboard": False},
            ):
                ray.init()
                worker_res_list = self.get_ray_worker_resources_list()
                assert len(worker_res_list) == max_worker_nodes
                for worker_res in worker_res_list:
                    assert (
                        worker_res["CPU"] == num_cpus_worker_node
                        and worker_res["memory"] == mem_worker_node
                        and worker_res["object_store_memory"]
                        == object_store_mem_worker_node
                    )

    def test_public_api(self):
        try:
            ray_temp_root_dir = tempfile.mkdtemp(dir="/tmp")
            collect_log_to_path = tempfile.mkdtemp(dir="/tmp")
            # Test the case that `collect_log_to_path` directory does not exist.
            shutil.rmtree(collect_log_to_path, ignore_errors=True)
            setup_ray_cluster(
                max_worker_nodes=MAX_NUM_WORKER_NODES,
                num_cpus_worker_node=1,
                num_gpus_worker_node=0,
                collect_log_to_path=collect_log_to_path,
                ray_temp_root_dir=ray_temp_root_dir,
                head_node_options={"include_dashboard": True},
            )

            assert (
                os.environ["RAY_ADDRESS"]
                == ray.util.spark.cluster_init._active_ray_cluster.address
            )

            ray.init()

            @ray.remote
            def f(x):
                return x * x

            futures = [f.remote(i) for i in range(32)]
            results = ray.get(futures)
            assert results == [i * i for i in range(32)]

            shutdown_ray_cluster()

            assert "RAY_ADDRESS" not in os.environ

            time.sleep(7)
            # assert temp dir is removed.
            assert len(os.listdir(ray_temp_root_dir)) == 1 and os.listdir(
                ray_temp_root_dir
            )[0].endswith(".lock")

            # assert logs are copied to specified path
            listed_items = os.listdir(collect_log_to_path)
            assert len(listed_items) == 1 and listed_items[0].startswith("ray-")
            log_dest_dir = os.path.join(
                collect_log_to_path, listed_items[0], socket.gethostname()
            )
            assert os.path.exists(log_dest_dir) and len(os.listdir(log_dest_dir)) > 0
        finally:
            if ray.util.spark.cluster_init._active_ray_cluster is not None:
                # if the test raised error and does not destroy cluster,
                # destroy it here.
                ray.util.spark.cluster_init._active_ray_cluster.shutdown()
                time.sleep(5)
            shutil.rmtree(ray_temp_root_dir, ignore_errors=True)
            shutil.rmtree(collect_log_to_path, ignore_errors=True)

    def test_autoscaling(self):
        for max_worker_nodes, num_cpus_worker_node, min_worker_nodes in [
            (self.max_spark_tasks, self.num_cpus_per_spark_task, 0),
            (self.max_spark_tasks // 2, self.num_cpus_per_spark_task * 2, 0),
            (self.max_spark_tasks, self.num_cpus_per_spark_task, 1),
        ]:
            num_ray_task_slots = self.max_spark_tasks // (
                num_cpus_worker_node // self.num_cpus_per_spark_task
            )
            (
                mem_worker_node,
                object_store_mem_worker_node,
                _,
            ) = _calc_mem_per_ray_worker_node(
                num_task_slots=num_ray_task_slots,
                physical_mem_bytes=_RAY_ON_SPARK_WORKER_PHYSICAL_MEMORY_BYTES,
                shared_mem_bytes=_RAY_ON_SPARK_WORKER_SHARED_MEMORY_BYTES,
                configured_heap_memory_bytes=None,
                configured_object_store_bytes=None,
            )

            with _setup_ray_cluster(
                max_worker_nodes=max_worker_nodes,
                min_worker_nodes=min_worker_nodes,
                num_cpus_worker_node=num_cpus_worker_node,
                num_gpus_worker_node=0,
                head_node_options={"include_dashboard": False},
                autoscale_idle_timeout_minutes=0.1,
            ):
                ray.init()
                worker_res_list = self.get_ray_worker_resources_list()
                assert len(worker_res_list) == min_worker_nodes

                @ray.remote(num_cpus=num_cpus_worker_node)
                def f(x):
                    import time

                    time.sleep(5)
                    return x * x

                # Test scale up
                futures = [f.remote(i) for i in range(8)]
                results = ray.get(futures)
                assert results == [i * i for i in range(8)]

                worker_res_list = self.get_ray_worker_resources_list()
                assert len(worker_res_list) == max_worker_nodes and all(
                    worker_res_list[i]["CPU"] == num_cpus_worker_node
                    and worker_res_list[i]["memory"] == mem_worker_node
                    and worker_res_list[i]["object_store_memory"]
                    == object_store_mem_worker_node
                    for i in range(max_worker_nodes)
                )

                # Test scale down
                wait_for_condition(
                    lambda: len(self.get_ray_worker_resources_list())
                    == min_worker_nodes,
                    timeout=60,
                    retry_interval_ms=1000,
                )
                if min_worker_nodes > 0:
                    # Test scaling down keeps nodes number >= min_worker_nodes
                    time.sleep(30)
                    assert len(self.get_ray_worker_resources_list()) == min_worker_nodes


class TestBasicSparkCluster(RayOnSparkCPUClusterTestBase):
    @classmethod
    def setup_class(cls):
        cls.num_total_cpus = 2
        cls.num_total_gpus = 0
        cls.num_cpus_per_spark_task = 1
        cls.num_gpus_per_spark_task = 0
        cls.max_spark_tasks = 2
        os.environ["SPARK_WORKER_CORES"] = "2"
        cls.spark = (
            SparkSession.builder.master("local-cluster[1, 2, 1024]")
            .config("spark.task.cpus", "1")
            .config("spark.task.maxFailures", "1")
            .config("spark.executorEnv.RAY_ON_SPARK_WORKER_CPU_CORES", "2")
            .getOrCreate()
        )


class TestSparkLocalCluster:
    @classmethod
    def setup_class(cls):
        cls.spark = (
            SparkSession.builder.master("local[2]")
            .config("spark.task.cpus", "1")
            .config("spark.task.maxFailures", "1")
            .getOrCreate()
        )

    @classmethod
    def teardown_class(cls):
        time.sleep(10)  # Wait all background spark job canceled.
        cls.spark.stop()

    def test_basic(self):
        local_addr, remote_addr = setup_ray_cluster(
            max_worker_nodes=2,
            head_node_options={"include_dashboard": False},
            collect_log_to_path="/tmp/ray_log_collect",
        )

        for cluster_addr in [local_addr, remote_addr]:
            ray.init(address=cluster_addr)

            @ray.remote
            def f(x):
                return x * x

            futures = [f.remote(i) for i in range(32)]
            results = ray.get(futures)
            assert results == [i * i for i in range(32)]

            ray.shutdown()

        shutdown_ray_cluster()

    def test_use_driver_resources(self):
        setup_ray_cluster(
            max_worker_nodes=1,
            num_cpus_head_node=3,
            num_gpus_head_node=2,
            object_store_memory_head_node=256 * 1024 * 1024,
            head_node_options={"include_dashboard": False},
            min_worker_nodes=0,
        )

        ray.init()
        head_resources_list = []
        for node in ray.nodes():
            if node["Alive"] and node["Resources"].get("CPU", 0) == 3:
                head_resources_list.append(node["Resources"])
        assert len(head_resources_list) == 1
        head_resources = head_resources_list[0]
        assert head_resources.get("GPU", 0) == 2

        shutdown_ray_cluster()

    def test_setup_global_ray_cluster(self):
        shutil.rmtree("/tmp/ray", ignore_errors=True)

        assert ray.util.spark.cluster_init._global_ray_cluster_cancel_event is None

        def start_serve_thread():
            def serve():
                try:
                    with mock.patch(
                        "ray.util.spark.cluster_init.get_spark_session",
                        return_value=self.spark,
                    ):
                        setup_global_ray_cluster(
                            max_worker_nodes=1,
                            min_worker_nodes=0,
                        )
                except BaseException:
                    # For debugging testing failure.
                    import traceback

                    traceback.print_exc()
                    raise

            threading.Thread(target=serve, daemon=True).start()

        start_serve_thread()

        wait_for_condition(
            (
                lambda: ray.util.spark.cluster_init._global_ray_cluster_cancel_event
                is not None
            ),
            timeout=120,
            retry_interval_ms=10000,
        )

        # assert it uses default temp directory
        assert os.path.exists("/tmp/ray")

        # assert we can connect to it on client server port 10001
        assert (
            ray.util.spark.cluster_init._active_ray_cluster.ray_client_server_port
            == 10001
        )

        with mock.patch("ray.util.spark.cluster_init._active_ray_cluster", None):
            # assert we cannot create another global mode cluster at a time
            with pytest.raises(
                ValueError,
                match=re.compile(
                    "Acquiring global lock failed for setting up new global mode "
                    "Ray on spark cluster"
                ),
            ):
                setup_global_ray_cluster(
                    max_worker_nodes=1,
                    min_worker_nodes=0,
                )

        # shut down the cluster
        ray.util.spark.cluster_init._global_ray_cluster_cancel_event.set()

        # assert temp directory is deleted
        wait_for_condition(
            lambda: not os.path.exists("/tmp/ray"),
            timeout=60,
            retry_interval_ms=10000,
        )

    def test_autoscaling_config_generation(self):
        from ray.util.spark.cluster_init import AutoscalingCluster

        autoscaling_cluster = AutoscalingCluster(
            head_resources={
                "CPU": 3,
                "GPU": 4,
                "memory": 10000000,
                "object_store_memory": 20000000,
            },
            worker_node_types={
                "ray.worker": {
                    "resources": {
                        "CPU": 5,
                        "GPU": 6,
                        "memory": 30000000,
                        "object_store_memory": 40000000,
                    },
                    "node_config": {},
                    "min_workers": 0,
                    "max_workers": 100,
                },
            },
            extra_provider_config={
                "extra_aa": "abc",
                "extra_bb": 789,
            },
            upscaling_speed=2.0,
            idle_timeout_minutes=3.0,
        )

        config = autoscaling_cluster._config

        assert config["max_workers"] == 100

        assert config["available_node_types"]["ray.head.default"] == {
            "resources": {
                "CPU": 3,
                "GPU": 4,
                "memory": 10000000,
                "object_store_memory": 20000000,
            },
            "node_config": {},
            "max_workers": 0,
        }
        assert config["available_node_types"]["ray.worker"] == {
            "resources": {
                "CPU": 5,
                "GPU": 6,
                "memory": 30000000,
                "object_store_memory": 40000000,
            },
            "node_config": {},
            "min_workers": 0,
            "max_workers": 100,
        }
        assert config["upscaling_speed"] == 2.0
        assert config["idle_timeout_minutes"] == 3.0
        assert config["provider"]["extra_aa"] == "abc"
        assert config["provider"]["extra_bb"] == 789

    def test_start_ray_node_in_new_process_group(self):
        from ray.util.spark.cluster_init import _start_ray_head_node

        proc, _ = _start_ray_head_node(
            [
                sys.executable,
                "-m",
                "ray.util.spark.start_ray_node",
                "--head",
                "--block",
                "--port=44335",
            ],
            synchronous=False,
            extra_env={
                "RAY_ON_SPARK_COLLECT_LOG_TO_PATH": "",
                "RAY_ON_SPARK_START_RAY_PARENT_PID": str(os.getpid()),
            },
        )
        time.sleep(10)

        # Assert the created Ray head node process has a different
        # group id from parent process group id.
        assert os.getpgid(proc.pid) != os.getpgrp()
        proc.terminate()


if __name__ == "__main__":
    sys.exit(pytest.main(["-sv", __file__]))
